#!/usr/bin/env python3
"""
数据库表信息查询脚本
使用MCP工具查询数据库中的表信息，包括表列表、表结构等详细信息
"""

import asyncio
import json
import sys
import os
from pathlib import Path

# 添加项目根目录到Python路径
project_root = Path(__file__).parent
sys.path.insert(0, str(project_root))

from src.mysql_mcp_server.config import DatabaseConfig
from src.mysql_mcp_server.mcp_server import create_mcp_server


async def query_table_info():
    """查询数据库表信息的主函数"""
    print("🚀 启动数据库表信息查询工具...")
    
    try:
        # 创建数据库配置，从环境变量加载
        config = DatabaseConfig.from_env()
        print(f"✅ 配置加载成功: {config.host}:{config.port}/{config.database}")
        
        # 创建MCP服务器
        mcp_server = await create_mcp_server(config)
        print("✅ MCP服务器创建成功")
        
        # 启动服务器（初始化数据库连接池）
        await mcp_server.start()
        print("✅ MCP服务器启动完成")
        
        # 初始化MCP协议（设置initialized标志）
        initialize_request = {
            "jsonrpc": "2.0",
            "id": 0,
            "method": "initialize",
            "params": {
                "protocolVersion": "2024-11-05",
                "clientInfo": {
                    "name": "table-query-client",
                    "version": "1.0.0"
                }
            }
        }
        
        response = await mcp_server.handle_request(json.dumps(initialize_request))
        init_result = json.loads(response)
        print("✅ MCP协议初始化完成")
        
        # 1. 获取所有表列表
        print("\n📋 获取数据库表列表...")
        print("=" * 50)
        
        tables_request = {
            'jsonrpc': '2.0',
            'id': 1,
            'method': 'tools/call',
            'params': {
                'name': 'list_tables',
                'arguments': {}
            }
        }
        
        response = await mcp_server.handle_request(json.dumps(tables_request))
        tables_result = json.loads(response)
        
        if tables_result and 'result' in tables_result and tables_result['result'] and 'content' in tables_result['result']:
            tables_info = json.loads(tables_result['result']['content'][0]['text'])
            
            if tables_info and 'tables' in tables_info:
                table_names = tables_info['tables']
                print(f"✅ 找到 {len(table_names)} 个表:")
                
                # 显示所有表名
                for i, table in enumerate(table_names, 1):
                    print(f"  {i:2d}. 📋 {table}")
                
                # 2. 查询指定表的详细结构
                print(f"\n🏗️  查询表结构详情...")
                print("=" * 50)
                
                # 选择前几个表进行详细查询
                sample_tables = table_names[:5]  # 查询前5个表
                
                for i, table_name in enumerate(sample_tables):
                    print(f"\n📝 表: {table_name}")
                    print("-" * 30)
                    
                    try:
                        # 查询表结构
                        describe_request = {
                            'jsonrpc': '2.0',
                            'id': i + 2,
                            'method': 'tools/call',
                            'params': {
                                'name': 'describe_table',
                                'arguments': {
                                    'table_name': table_name
                                }
                            }
                        }
                        
                        response = await mcp_server.handle_request(json.dumps(describe_request))
                        describe_result = json.loads(response)
                        
                        if describe_result and 'result' in describe_result and describe_result['result'] and 'content' in describe_result['result']:
                            text_content = describe_result['result']['content'][0]['text']
                            try:
                                describe_info = json.loads(text_content)
                            except json.JSONDecodeError as e:
                                print(f"   ❌ JSON解析错误: {e}")
                                print(f"   原始内容: {text_content[:200]}...")
                                continue
                            
                            # 检查是否有错误信息
                            if "error" in describe_info:
                                error_msg = describe_info['error']
                                error_type = describe_info.get('error_type', '未知错误类型')
                                print(f"   ❌ {error_msg} (类型: {error_type})")
                                continue
                            
                            if describe_info and 'columns' in describe_info:
                                columns = describe_info['columns']
                                print(f"   字段数量: {len(columns)}")
                                print("   字段详情:")
                                
                                for col in columns:
                                    key_info = ""
                                    if col.get('Key') == 'PRI':
                                        key_info = " 🔑(主键)"
                                    elif col.get('Key') == 'UNI':
                                        key_info = " 🔒(唯一)"
                                    elif col.get('Key') == 'MUL':
                                        key_info = " 🔗(索引)"
                                    
                                    null_info = "允许NULL" if col.get('Null') == 'YES' else "不允许NULL"
                                    default_val = col.get('Default', 'N/A')
                                    
                                    print(f"     • {col['Field']:<20} {col['Type']:<15}{key_info}")
                                    print(f"       {null_info}, 默认值: {default_val}")
                            else:
                                print("   ❌ 无法获取表结构")
                        elif describe_result.get("error"):
                            print(f"   ❌ MCP错误: {describe_result['error'].get('message', '未知错误')}")
                        else:
                            print("   ❌ 表结构查询失败")
                    except Exception as e:
                        print(f"   ❌ 查询异常: {str(e)}")
                        import traceback
                        print(f"   详细错误: {traceback.format_exc()}")
                
                if len(table_names) > 5:
                    print(f"\n💡 提示: 只显示了前5个表的结构，总共有 {len(table_names)} 个表")
                
                # 3. 执行一些统计查询
                print(f"\n📊 数据库统计信息...")
                print("=" * 50)
                
                # 查询表数量
                count_query = "SELECT COUNT(*) as table_count FROM information_schema.tables WHERE table_schema = DATABASE()"
                count_request = {
                    'jsonrpc': '2.0',
                    'id': 3,
                    'method': 'tools/call',
                    'params': {
                        'name': 'execute_query',
                        'arguments': {
                            'query': count_query,
                            'limit': 1
                        }
                    }
                }
                
                response = await mcp_server.handle_request(json.dumps(count_request))
                count_result = json.loads(response)
                
                if count_result and 'result' in count_result and count_result['result'] and 'content' in count_result['result']:
                    count_info = json.loads(count_result['result']['content'][0]['text'])
                    
                    if count_info and 'results' in count_info:
                        table_count = count_info['results'][0]['table_count']
                        print(f"📋 数据库表总数: {table_count}")
                
                # 查询数据库大小信息
                size_query = """
                SELECT 
                    table_schema as '数据库',
                    COUNT(*) as '表数量',
                    ROUND(SUM(data_length + index_length) / 1024 / 1024, 2) as '大小(MB)'
                FROM information_schema.tables 
                WHERE table_schema = DATABASE()
                GROUP BY table_schema
                """
                
                size_request = {
                    'jsonrpc': '2.0',
                    'id': 4,
                    'method': 'tools/call',
                    'params': {
                        'name': 'execute_query',
                        'arguments': {
                            'query': size_query,
                            'limit': 1
                        }
                    }
                }
                
                response = await mcp_server.handle_request(json.dumps(size_request))
                size_result = json.loads(response)
                
                if size_result and 'result' in size_result and size_result['result'] and 'content' in size_result['result']:
                    size_info = json.loads(size_result['result']['content'][0]['text'])
                    
                    if size_info and 'results' in size_info:
                        db_info = size_info['results'][0]
                        print(f"💾 数据库: {db_info['数据库']}")
                        print(f"📊 表数量: {db_info['表数量']}")
                        print(f"📏 大小: {db_info['大小(MB)']} MB")
                
            else:
                print("❌ 无法获取表列表")
        else:
            print("❌ 表列表查询失败")
            
    except Exception as e:
        print(f"❌ 查询过程中发生错误: {e}")
    
    finally:
        # 停止MCP服务器
        await mcp_server.stop()
        print("\n🎉 查询完成！MCP服务器已停止")


def main():
    """主入口函数"""
    print("=" * 60)
    print("🔍 数据库表信息查询工具")
    print("=" * 60)
    
    try:
        asyncio.run(query_table_info())
    except KeyboardInterrupt:
        print("\n⚠️  用户中断操作")
    except Exception as e:
        print(f"\n❌ 程序执行失败: {e}")


if __name__ == "__main__":
    main()